import tkinter as tk
import time
import threading
import random
import chess
import PIL as pl
import sys
from PIL import ImageGrab as IG
from ChessDQN import ChessDQN
import numpy as np
import tensorflow as tf
import time

from win32 import win32api, win32gui, win32print
from win32.lib import win32con
from win32.win32api import GetSystemMetrics

class TTTGym:
    def __init__(self):
        self._INIT_CHESS = [0, 0, 0, 0, 0, 0, 0, 0, 0]  # 原始棋盘
        self._the_chess_status = [0, 0, 0, 0, 0, 0, 0, 0, 0]  # 记录棋盘
        self._pre_chess_status = [0, 0, 0, 0, 0, 0, 0, 0, 0] # 上一步棋盘记录下棋的位置

        # 创建棋盘UI
        self.ui_main_ttt = tk.Tk()
        self.ui_main_ttt.title('井字棋 -> Fighting')
        self.ui_main_ttt.geometry("300x300")
        self.ui_main_ttt.resizable()
        self.ui_string_tips = tk.StringVar(self.ui_main_ttt)

        self.ui_chess_pos = [tk.StringVar(self.ui_main_ttt) for _ in range(9)]

        # self.ui_layout_frame_top = tk.Frame(self.ui_main_ttt)
        self.ui_layout_frame_cont = tk.Frame(self.ui_main_ttt)
        self.ui_layout_frame_bot = tk.Frame(self.ui_main_ttt)
        self.ui_layout_frame_cont1 = tk.Frame(self.ui_layout_frame_cont)
        self.ui_layout_frame_cont2 = tk.Frame(self.ui_layout_frame_cont)
        self.ui_layout_frame_cont3 = tk.Frame(self.ui_layout_frame_cont)

        # 棋盘显示label 0~9
        self.ui_chess_albel0 = tk.Label(self.ui_layout_frame_cont1, textvariable=self.ui_chess_pos[0], font=("幼圆", 30), padx=0)
        self.ui_chess_albel1 = tk.Label(self.ui_layout_frame_cont1, textvariable=self.ui_chess_pos[1], font=("幼圆", 30), padx=0)
        self.ui_chess_albel2 = tk.Label(self.ui_layout_frame_cont1, textvariable=self.ui_chess_pos[2], font=("幼圆", 30), padx=0)

        self.ui_chess_albel3 = tk.Label(self.ui_layout_frame_cont2, textvariable=self.ui_chess_pos[3], font=("幼圆", 30), padx=0)
        self.ui_chess_albel4 = tk.Label(self.ui_layout_frame_cont2, textvariable=self.ui_chess_pos[4], font=("幼圆", 30), padx=0)
        self.ui_chess_albel5 = tk.Label(self.ui_layout_frame_cont2, textvariable=self.ui_chess_pos[5], font=("幼圆", 30), padx=0)

        self.ui_chess_albel6 = tk.Label(self.ui_layout_frame_cont3, textvariable=self.ui_chess_pos[6], font=("幼圆", 30), padx=0)
        self.ui_chess_albel7 = tk.Label(self.ui_layout_frame_cont3, textvariable=self.ui_chess_pos[7], font=("幼圆", 30), padx=0)
        self.ui_chess_albel8 = tk.Label(self.ui_layout_frame_cont3, textvariable=self.ui_chess_pos[8], font=("幼圆", 30), padx=0)

        self.ui_top_label = tk.Label(self.ui_layout_frame_bot, justify=tk.CENTER, textvariable=self.ui_string_tips, font=("幼圆", 20), padx=0)

        self.ui_chess_albel0.pack(side=tk.LEFT)
        self.ui_chess_albel1.pack(side=tk.LEFT)
        self.ui_chess_albel2.pack(side=tk.LEFT)
        self.ui_chess_albel3.pack(side=tk.LEFT)
        self.ui_chess_albel4.pack(side=tk.LEFT)
        self.ui_chess_albel5.pack(side=tk.LEFT)
        self.ui_chess_albel6.pack(side=tk.LEFT)
        self.ui_chess_albel7.pack(side=tk.LEFT)
        self.ui_chess_albel8.pack(side=tk.LEFT)

        self.ui_layout_frame_cont1.pack()
        self.ui_layout_frame_cont2.pack()
        self.ui_layout_frame_cont3.pack()
        self.ui_layout_frame_cont.pack()
        self.ui_layout_frame_bot.pack()

    def reset(self):
        '''
        重置游戏状态 而这个状态就是强化学习的输入数据，即train_data
        :return: 返回当前的状态，是一个矩阵信息
        '''

        self._the_chess_status = self._INIT_CHESS.copy()
        self._pre_chess_status = self._INIT_CHESS.copy()

        random.seed(time.time())

        # 随机下棋
        if random.randint(0, 1):
            self._the_chess_status[random.randint(0, 8)] = -1

        self.ui_main_ttt.update()

        # return self._the_chess_status
        return self._the_chess_status

    def resetRState(self):
        '''
                重置游戏状态 而这个状态就是强化学习的输入数据，即train_data
                :return: 返回当前的状态，是一个矩阵信息
                '''

        self._the_chess_status = self._INIT_CHESS.copy()
        self._pre_chess_status = self._INIT_CHESS.copy()

        random.seed(time.time())

        # 随机下棋
        if random.randint(0, 1):
            self._the_chess_status[random.randint(0, 8)] = -1

        self.ui_main_ttt.update()
        cur_chess_img = self.__get_cur_chess(self.__get_cur_scle())

        return cur_chess_img

    def render(self):
        '''
        绘制当前的游戏状态
        1 表示AI -1表示传统AI
        :return:
        '''

        for pos, ui_pos in zip(self._the_chess_status, self.ui_chess_pos):
            if pos == 1:
                ui_pos.set('|X|')
            elif pos == -1:
                ui_pos.set('|O|')
            else:
                ui_pos.set('| |')

        self.ui_main_ttt.update()

    def loop(self):
        '''
        tkinter图形界面通过mainloop进入消息循环显示界面
        todo 但是由于tkinter界面ui存在显示与捕获的问题，暂时不做
        :return:
        '''
        self.ui_main_ttt.mainloop()

    def __get_cur_scle(self):
        if sys.platform.startswith('win32'):
            """获取真实的分辨率"""
            hDC = win32gui.GetDC(0)
            # 横向分辨率
            # real_w = win32print.GetDeviceCaps(hDC, win32con.DESKTOPHORZRES)
            # 纵向分辨率
            real_h = win32print.GetDeviceCaps(hDC, win32con.DESKTOPVERTRES)

            # logit_w = GetSystemMetrics(0)
            logit_h = GetSystemMetrics(1)

            return round(real_h / logit_h, 2)
        else:
            return 2

    def __get_cur_chess(self, scale=2):
        # 截图
        grab_pos = (self.ui_layout_frame_cont.winfo_rootx() + 1, self.ui_layout_frame_cont.winfo_rooty() + 1,
                    self.ui_layout_frame_cont.winfo_width() + self.ui_layout_frame_cont.winfo_rootx() - 1,
                    self.ui_layout_frame_cont.winfo_height() + self.ui_layout_frame_cont.winfo_rooty() - 1)

        if grab_pos[2] > grab_pos[0] and grab_pos[3] > grab_pos[1]:
            pic = IG.grab([scale * pos for pos in grab_pos])
            pic.save("./data/cur_chess.bmp")
        return np.array(pic)

    def step(self, action):
        '''
        进行下一步动作返回动作的反馈
        :param action: 下一步的动作
        :return: 观察、奖励、是否结束、信息(返回额外的信息数据)
        '''
        """
            创建一个机器对打局
            i 表示当前是第几轮下棋
            id表示当前是哪个ai先下
            """

        # cur_chess_pos = self._INIT_CHESS[:]
        cur_chess_action = action[:]
        cur_chess_status = self._the_chess_status
        pre_chess_status = cur_chess_status

        if cur_chess_status[np.argmax(cur_chess_action)] != 0:
            return cur_chess_status, -10, -1, np.argmax(action), pre_chess_status

        cur_chess_status[np.argmax(cur_chess_action)] = 1
        id = -1

        reward = 1
        is_done = 0
        # cur_chess_img = self.__get_cur_chess(self.__get_cur_scle())
        # 如果棋盘依旧还剩格子未下
        # todo 增加判断棋局是否已经结束，如果已经结束则直接返回

        if chess.count_zero(cur_chess_status) > 0:
            # 当前id ai下棋
            pos = chess.get_next_o(cur_chess_status, id)
            if pos != id:
                # todo 这里可以记录下棋的位置
                cur_chess_status[int(pos)] = id
                # cur_chess_pos[int(pos)] = id
            else:
                # 如果找不到下棋的位置，那就随意找一个空格位下棋
                # todo 这里可以记录下棋的位置
                for xx in range(9):
                    if cur_chess_status[xx] == 0:
                        cur_chess_status[xx] = id
                        # cur_chess_pos[xx] = id

            reward = 1
            is_done = 0
            if chess.is_win(cur_chess_status, -1) == -1:
                name = "X"
                reward = -10
                is_done = -1
                print("第 {} 局 ： {} 赢了！".format(i + 1, name) + ' ' + str(new_chess))
                tips.set("第 {} 局 ： {} 赢了！".format(i + 1, name))
                count_x = count_x + 1
            elif chess.is_win(cur_chess_status, 1) == 1:
                name = "O"
                reward = 10
                is_done = 10
                print("第 {} 局 ： {} 赢了！".format(i + 1, name) + ' ' + str(new_chess))
                tips.set("第 {} 局 ： {} 赢了！".format(i + 1, name))
                count_x = count_x + 1
            elif chess.count_zero(cur_chess_status) == 0:
                is_done = 1
                print("第 {} 局 ： 平局".format(i + 1) + ' ' + str(cur_chess_status))
                tips.set("第 {} 局 ： 平局".format(i + 1))
            else:
                pass

        self._the_chess_status = cur_chess_status
        # return cur_chess_img, reward, is_done, cur_chess_pos
        return cur_chess_status, reward, is_done, np.argmax(action), pre_chess_status